from numpy import asarray,where,
from pylab import plot,show,legend

class InterpCubicSpline:
    """
    Given a set of points x,y, interpolates a cubic spline that passes
    through all points.
    
    Instantiate the class with two arrays of points: x and y =
    f(x). Loosely based on the Numerical recipes routines.

    Inputs:

    x                  : array of x values
    y                  : array of y values (y = f(x))
    first_deriv = None : derivative of f(x) at x[0]
    last_deriv  = None : derivative of f(x) at x[-1]

    After initialisation, the instance can be called with an array of
    values xp, and will return the cubic-spline interpolated values yp
    = f(xp).

    The spline can be re-calculated with a new first and last
    derivative using the set_d2() method.

    If you want to calculate a new spline using a different set of x
    values, you'll have to instantiate a new class.
    """
    def __init__(self,x,y,first_deriv=None,last_deriv=None):
        x = np.asarray(x).astype('float')
        lx = list(x)
        # check all x values are unique
        if any((lx.count(val) > 1) for val in set(lx)):
            raise Exception('non-unique x values were found!')
        y = np.asarray(y).astype('float')
        cond = x.argsort()                    # sort arrays
        self.x = x[cond]
        self.y = y[cond]
        self.npts = len(x)
        self.set_d2(first_deriv,last_deriv)
        
    def __call__(self,xp):
        """ Given an array of x values, returns cubic-spline
        interpolated values yp = f(xp) using the derivatives
        calculated in set_d2().
        """
        x = self.x;  y = self.y;  npts = self.npts;  d2 = self.d2

        # make xp into an array
        if not hasattr(xp,'__len__'):  xp = xp,
        xp = np.asarray(xp)

        # for each xp value, find the closest x value above and below
        i2 = np.searchsorted(x,xp)

        # account for xp values outside x range
        i2 = np.where(i2 == npts, npts-1, i2)
        i2 = np.where(i2 == 0, 1, i2)
        i1 = i2 - 1

        h = x[i2] - x[i1]
        a = (x[i2] - xp) / h
        b = (xp - x[i1]) / h
        temp = (a**3 - a)*d2[i1] +  (b**3 - b)*d2[i2]
        yp = a * y[i1] + b * y[i2] + temp * h*h/6.

        return yp

    def tri_diag(self,temp,d2):
        x = self.x;  y = self.y;  npts = self.npts        
        for i in range(1,npts-1):
            ratio = (x[i]-x[i-1]) / (x[i+1]-x[i-1])
            denom = ratio * d2[i-1] + 2.       # 2 if x vals equally spaced
            d2[i] = (ratio - 1.) / denom       # -0.5 if x vals equally spaced
            temp[i] = (y[i+1]-y[i])/(x[i+1]-x[i]) - (y[i]-y[i-1])/(x[i]-x[i-1])
            temp[i] = (6.*temp[i]/(x[i+1]-x[i-1]) - ratio * temp[i-1]) / denom
        return temp
            
    def set_d2(self, first_deriv=None, last_deriv=None,verbose=False):
        """ Calculates the second derivative of a cubic spline
        function y = f(x) for each value in array x. This is called once by
        __init__() when a new class instance is created.

        optional inputs:

        first_deriv = None : 1st derivative of f(x) at x[0].  If None,
                             then 2nd derivative is set to 0 ('natural').
        last_deriv  = None : 1st derivative of f(x) at x[-1].  If None,
                             then 2nd derivative is set to 0 ('natural').
        """
        if verbose:  print 'first deriv,last deriv',first_deriv,last_deriv
        x = self.x;  y = self.y;  npts = self.npts
        d2 = np.empty(npts)
        temp = np.empty(npts-1)

        if first_deriv is None:
            if verbose:  print "Lower boundary condition set to 'natural'"
            d2[0] = 0.
            temp[0] = 0.
        else:
            d2[0] = -0.5
            temp[0] = 3./(x[1]-x[0]) * ((y[1]-y[0])/(x[1]-x[0]) - first_deriv)
            #print temp[0]

        temp = self.tri_diag(temp,d2)

        if last_deriv is None:
            if verbose:  print "Upper boundary condition set to 'natural'"
            qn = 0.
            un = 0.
        else:
            qn = 0.5
            #print (y[-1]-y[-2])/(x[-1]-x[-2])
            un = 3./(x[-1]-x[-2]) * (last_deriv - (y[-1]-y[-2])/(x[-1]-x[-2]))
            #print 'un=',un

        d2[-1] = (un - qn*temp[-1]) / (qn*d2[-2] + 1.)
        for i in reversed(range(npts-1)):
            d2[i] = d2[i] * d2[i+1] + temp[i]

        self.d2 = d2

if 0:
    npts = 5

    #x = np.arange(npts,dtype='float')
    x = [-6,-5,0,0.1,3,4,7,25]
    #x = x + np.random.randn(npts)
    #y = np.random.randn(npts)
    y = [5,10,2,0,-0.5,6,9,-1]#,4,7,4,3,-6,3,1,-3,-1,-4,2,7,9,2,1]
    #def fy(i): return y[int(i)]
    xvals = np.linspace(-6.5,25,20000)
    spl = InterpCubicSpline(x,y)#,last_deriv=1,first_deriv=0)
    spl1 = InterpCubicSpline(x,y,first_deriv=0,last_deriv=0)
    #spl2 = Interpolator('junk',fy,x)
    #spl3 = Interpolator('junk',fy,x)
    pl.axis('equal')
    plot(x,y,'rx',mew=2,ms=8,label=None)
    plot(xvals,spl(xvals),label='natural spline')
    plot(xvals,spl1(xvals),label='end derivs')
    #plot(xvals,[spl3(val) for val in xvals],label='cookbook 2')
    plot([2,6],[-3,1],':k')
    #un = UnivariateSpline(x,y,s=0)
    #plot(xvals,un(xvals),label='scipy spline')
    #legend()
